import  os
os.environ["MXNET_CUDNN_AUTOTUNE_DEFAULT"] = "0"
from gluoncv import model_zoo, data, utils
from matplotlib import pyplot as plt
from lib.common import lsdir
import mxnet as mx
import gluoncv
print(mx.__version__,gluoncv.__version__)
net = model_zoo.get_model('faster_rcnn_resnet50_v2a_coco', pretrained=False,classes = ["head","helmet"])
net.load_parameters("output/no_latex/faster_rcnn_resnet50_v1b_custom_0017_0.8273.params",allow_missing=True,ignore_extra=True)
files = lsdir("/data1/zyx/yks/dataset/ocr_formula/val",suffix=".png")
net.collect_params().reset_ctx([mx.gpu(0)])
import time
for im_fname in files:   
    t0= time.time()
    x, img = data.transforms.presets.rcnn.load_test(im_fname, short=600,max_size=1024)
    print('Shape of pre-processed image:', x.shape)
    x = x.as_in_context(mx.gpu(0))
    class_IDs, scores, bounding_boxs,_ = net(x)
    print(time.time()-t0)
    ax = utils.viz.plot_bbox(img, bounding_boxs[0], scores[0],
                             class_IDs[0], class_names=net.classes)
    plt.show()